#!/bin/bash
#
# a simple initramfs generator inspired by Arch's mkinitcpio
# but wrote it simpler
#

PATH=/bin:/usr/bin:/sbin:/usr/sbin

add_binary() {
	local binpath binlib
	binpath=$(type -p $1)
	if [ -z "$binpath" ]; then
		msgwarn "missing binary: $1"
		return
	fi
	add_file $binpath
	if [ -L "$binpath" ]; then
		add_binary $(readlink -f $binpath)
		return
	fi
	case "$(file -bi $binpath)" in
		*application/x-sharedlib* | *application/x-executable* | *application/x-pie-executable*)
			binlib=$(ldd $binpath | grep -v "not" | sed "s/\t//" | cut -d " " -f1)
			for lib in $binlib; do
				case $lib in
					linux-vdso.so.1|linux-gate.so.1) continue ;;
				esac
				add_file $(PATH=/lib:/lib64:/usr/lib:/usr/lib64 type -p $lib)
				unset lib
			done;;
	esac
}

add_module() {
	local modname modpath
	if modinfo -k $KERNEL $1 &>/dev/null; then
		modname=$(modinfo -k $KERNEL -F name $1 | cut -d ' ' -f1 | head -n1)
		[ "$modname" = "name:" ] && return 0
		modpath=$(modinfo -k $KERNEL -F filename $1 | cut -d ' ' -f1 | head -n1)
		[ "$modpath" = "name:" ] && return 0
	else
		msg "missing module: $1"
		return
	fi
	[ -f $INITDIR/lib/modules/$KERNEL/kernel/$(basename $modpath) ] && return
	add_file "$modpath" lib/modules/$KERNEL/kernel/$(basename $modpath)
	modinfo -F firmware -k $KERNEL $modname | while read -r line; do
		if [ ! -f /lib/firmware/$line ]; then
			msgwarn "missing firmware for $modname: $line"
		else
			add_file /lib/firmware/$line
		fi
	done
	modinfo -F depends -k $KERNEL $modname | while IFS=',' read -r -a line; do
		for l in ${line[@]}; do
			add_module "$l"
		done
	done
}

add_file() {
	local src dest mode
	
	[ "$1" ] || return
	
	src=$1
	
	if [ ! -f "$src" ]; then
		msgwarn "file not found: $src"
		return
	fi
	
	if [ "${src:0:1}" != "/" ]; then
		msgwarn "absolute source path needed: $src"
		return
	fi
	
	if [ -z "$2" ]; then
		dest=${src/\//}
	else
		dest="$2"
	fi
	
	if [ "${dest:0:1}" = "/" ]; then
		msgwarn "destination path must without leading '/': $dest"
		return
	fi
	
	mode=${3:-$(stat -c %a "$src")}
	if [ -z "$mode" ]; then
		msgwarn "failed get file mode: $src"
		return
	fi
	
	install -Dm$mode $src $INITDIR/$dest
}

add_dir() {
	local path=$1 mode=${2:-755}
	
	if [[ -z $path ]] || [[ $path != /?* ]]; then
		return 1
	fi
	
	if [[ -d $INITDIR$path ]]; then
		return 0
	fi
	
	install -dm$mode $INITDIR$path
}

add_symlink() {
	local name=$1 target=$2
	
	[ "$name" ] || return 1
	
	if [ ! "$target" ]; then
		target=$(readlink -f $name)
		if [ ! "$target" ]; then
			msgerr "invalid symlink: $name"
			return 1
		fi
	fi
	
	add_dir "${name%/*}"
	
	if [ -L "$INITDIR$name" ]; then
		msgwarn "overwriting symlink: $name"
	fi
	
	ln -sfn "$target" $INITDIR$name
}

finalize_modules() {
	local file
	[ -d $INITDIR/lib/modules/$KERNEL/kernel ] || return
	for file in /lib/modules/$KERNEL/modules.*; do
		add_file $file
	done
	awk -F'/' '{ print "kernel/" $NF }' /lib/modules/$KERNEL/modules.order > $INITDIR/lib/modules/$KERNEL/modules.order
	depmod -b $INITDIR $KERNEL
}

run_build_hook() {
	if [ "${#HOOKS[@]}" -gt 0 ]; then
		for hook in ${HOOKS[@]}; do
			[ "$(echo ${DONEHOOK[@]} | tr ' ' '\n' | grep -x $hook)" = "$hook" ] && continue
			hookpath=$(PATH=$HOOKDIRS type -p $hook.hook)
			if [ "$hookpath" ]; then
				source "$hookpath"
				if [ "$(type -t build_hook)" = "function" ]; then
					msg "running build_hook: $hook"
					build_hook
				fi
				echo $hook >> $INITDIR/hook/hook.order
				add_file "$hookpath" hook/$hook 755
				DONEHOOK+=($hook)
				unset hookpath build_hook hook
			else
				msgwarn "missing hook: $hook"
			fi
		done
	fi
}

msg() {
	[ "$QUIET" ] && return
	echo ":: $@"
}

msgerr() {
	[ "$QUIET" ] && return
	echo "ERROR: $@"
}

msgwarn() {
	[ "$QUIET" ] && return
	echo "WARNING: $@"
}

cleanup() {
	rm -fr $INITDIR
}

interrupted() {
	cleanup
	exit 1
}

usage() {	
	cat << EOF
Usage:
  $(basename $0) [option] [argument]
  
Options:
  -k <version>  custom kernel version (default: $KERNEL)
  -o <output>   custom output name (default: $INITRAMFS)
  -i <init>     custom init file (default: $INITIN)
  -m <modules>  add extra modules (comma separated)
  -b <binaries  add extra binary (comma separated)
  -f <file>     add extra file (comma separated & absolute path)
  -c <config>   use custom config (default: $CONFIG)
  -A <hook>     add extra hook (comma separated, precedence over -a, -s & HOOKS)
  -a <hook>     add extra hook (comma separated, precedence over -s & after HOOKS)
  -s <hook>     skip hook defined in HOOKS (comma separated)
  -q            quiet mode
  -h            print this help msg
	
EOF
}

needarg() {
	if [ ! "$1" ]; then
		echo "ERROR: argument is needed for this option!"
		exit 1
	fi
}		

parse_opt() {
	while [ $1 ]; do
		case $1 in
			-k)	needarg $2
				KERNEL=$2
				shift 1 ;;
			-o)	needarg $2
				OUTPUT=$2
				shift 1 ;;
			-i)	needarg $2
				INITIN=$2
				shift 1 ;;
			-c)	needarg $2
				CONFIG=$2
				shift 1 ;;
			-A)	needarg $2
			    IFS=, read -r -a ADDEARLYHOOKS <<< $2
				shift 1 ;;
			-a)	needarg $2
			    IFS=, read -r -a ADDHOOKS <<< $2
				shift 1 ;;
			-s)	needarg $2
			    IFS=, read -r -a SKIPHOOKS <<< $2
				shift 1 ;;
			-m) needarg $2
			    IFS=, read -r -a ADDMODULES <<< $2
				shift 1 ;;
			-b) needarg $2
			    IFS=, read -r -a ADDBINARIES <<< $2
				shift 1 ;;
			-f) needarg $2
			    IFS=, read -r -a ADDFILES <<< $2
				shift 1 ;;
			-q)	QUIET=1 ;;
			-h)	usage; exit 0 ;;
			*)	echo "ERROR: invalid option '$1'"
				exit 1 ;;
		esac
		shift
	done
}

main() {
	parse_opt $@

	if [ "$UID" != "0" ]; then
		msgerr "need root access!"
		exit 1
	fi
	
	if [ ! $(type -p cpio) ]; then
		msgerr "'bsdcpio' not found, please install 'libarchive' package."
		exit 1
	fi
	
	if [ -f "$CONFIG" ]; then
		. "$CONFIG"
	else
		msgerr "config file '$CONFIG' not exist."
		exit 1
	fi
	
	if [ ! -d /lib/modules/"$KERNEL" ]; then
		msgerr "kernel directory '/lib/modules/$KERNEL' not exist."
		exit 1
	fi
	
	if [ ! -f "$INITIN" ]; then
		msgerr "init file '$INITIN' not exist."
		exit 1
	fi
	
	if [ "$OUTPUT" ]; then
		if [ $(basename $OUTPUT) != "$OUTPUT" ] && [ ! -d $(dirname $OUTPUT) ]; then
			msgerr "directory '$(dirname $OUTPUT)' for output '$(basename $OUTPUT)' not exist."
			exit 1
		elif [ -d "$OUTPUT" ]; then
			msgerr "'$OUTPUT' is a directory."
			exit 1
		fi
		INITRAMFS="$OUTPUT"
	fi
	
	# filter out skip hooks (-s)
	if [ "${#SKIPHOOKS[@]}" -gt 0 ] && [ "${#HOOKS[@]}" -gt 0 ]; then
		for s in ${!SKIPHOOKS[@]}; do
			for h in ${!HOOKS[@]}; do
				if [ "${SKIPHOOKS[s]}" = "${HOOKS[h]}" ]; then
					unset 'HOOKS[h]'
					break
				fi
			done
		done
	fi
	
	# add extra hooks (-a)
	if [ "${#ADDHOOKS[@]}" -gt 0 ]; then
		HOOKS+=(${ADDHOOKS[@]})
	fi
	
	# add extra early hooks (-A)
	if [ "${#ADDEARLYHOOKS[@]}" -gt 0 ]; then
		ADDEARLYHOOKS+=(${HOOKS[@]})
		HOOKS=(${ADDEARLYHOOKS[@]})
	fi
	
	# add extra modules (-m)
	if [ "${#ADDMODULES[@]}" -gt 0 ]; then
		MODULES+=(${ADDMODULES[@]})
	fi
	
	# add extra files (-f)
	if [ "${#ADDFILES[@]}" -gt 0 ]; then
		FILES+=(${ADDFILES[@]})
	fi
	
	# add extra binary (-b)
	if [ "${#ADDBINARIES[@]}" -gt 0 ]; then
		BINARIES+=(${ADDBINARIES[@]})
	fi

	[ "$QUIET" ] || echo "Generating initramfs..."
	
	mkdir -p $INITDIR/{hook,newroot}
	install -m0755 $INITIN $INITDIR/init
	run_build_hook
	
	if [ "${#BINARIES[@]}" -gt 0 ]; then
		msg "adding extra binaries..."
		for b in ${BINARIES[@]}; do
			add_binary "$b"
		done
	fi
	
	if [ "${#MODULES[@]}" -gt 0 ]; then
		msg "adding extra modules..."
		for m in ${MODULES[@]}; do
			add_module "$m"
		done
	fi
	
	if [ "${#FILES[@]}" -gt 0 ]; then
		msg "adding extra files..."
		for f in ${FILES[@]}; do
			add_file "$f"
		done
	fi
	
	finalize_modules

	msg "generating initramfs..."
	rm -f "$INITRAMFS"
	( cd $INITDIR ; find . | LANG=C cpio -o -H newc --quiet | gzip -9 ) > $INITRAMFS

	cleanup
	[ "$QUIET" ] || echo "Generating initramfs done: $INITRAMFS ($(du -h $INITRAMFS | awk '{print $1}'))"
	
	exit 0
}

trap "interrupted" SIGHUP SIGINT SIGQUIT SIGTERM

INITDIR="/tmp/mkinitramfs.$$"
KERNEL="$(uname -r)"
INITIN="/usr/share/mkinitramfs/init.in"
INITRAMFS="initrd-$KERNEL.img"
CONFIG="/etc/mkinitramfs.conf"
HOOKDIRS="/etc/mkinitramfs.d:/usr/share/mkinitramfs/hooks"

main $@
